from config import Opt

from models import Model

import torch
from transformers import AutoTokenizer
from utils import *
import copy
import re

SEP = '[SEP]'

VISUALIZATION_DICT = ['success','primary','info']

COLORS = [(255,0,255),(0,255,255),(255,255,0),(255,0,255),(0,255,255),(255,255,0)]

IGNORE_CHAR_SET = {'[SEP]'}

class Visualization:
    def __init__(self,opt,model):
        self.opt = opt
        self.model = model
        self.tokenizer = AutoTokenizer.from_pretrained(self.opt.bert_path)

        self.opt.ent_label_matrix = torch.tensor(self.opt.ent_label_matrix,dtype=torch.float32)
        self.opt.label_matrix = torch.tensor(self.opt.label_matrix,dtype=torch.float32)
        if torch.cuda.is_available() and self.opt.gpu >= 0:
            self.opt.ent_label_matrix = self.opt.ent_label_matrix.cuda(self.opt.gpu)
            self.opt.label_matrix = self.opt.label_matrix.cuda(self.opt.gpu)
            self.model = self.model.cuda(self.opt.gpu)
        self.use_longformer = 'Longformer' in self.opt.bert_path
        self.model.eval()

    def predict(self,dic):
        """
            data_dict : {'性别','年龄','主诉','现病史','既往史',...}
        """

        if '主诉' not in dic["主诉"]:
            dic['主诉'] = '主诉：'+dic['主诉']
        if '现病史' not in dic["现病史"]:
            dic['现病史'] = '现病史：'+dic['现病史']
        if '既往史' not in dic["既往史"]:
            dic['既往史'] = '既往史：'+dic['既往史']
        chief_complaint = '性别：' + dic['性别'] + '；年龄：'+get_age(dic['年龄']) + '；'+ dic["主诉"]
        now_history, past_history = dic["现病史"], dic["既往史"]

        doc = chief_complaint + '[SEP]' + now_history + '[SEP]' + past_history

        chief_complaint_tokens = self.tokenizer.tokenize(chief_complaint[:50]) + [SEP]
        now_history_tokens = self.tokenizer.tokenize(now_history[:500]) + [SEP]
        past_history_tokens = self.tokenizer.tokenize(past_history[:300]) + [SEP]
        doc_tokens = self.tokenizer.tokenize(doc[:850]) + [SEP]

        chief_complaint = self.tokenizer(chief_complaint[:50],return_tensors='pt')
        now_history = self.tokenizer(now_history[:500],return_tensors='pt')
        past_history = self.tokenizer(past_history[:300],return_tensors='pt')
        doc = self.tokenizer(doc[:850],return_tensors='pt')

        print(len(chief_complaint_tokens),chief_complaint['input_ids'].size(1)-1)
        print(len(now_history_tokens),now_history['input_ids'].size(1)-1)
        print(len(past_history_tokens),past_history['input_ids'].size(1)-1)
        print(len(doc_tokens),doc['input_ids'].size(1)-1)
        
        assert len(chief_complaint_tokens) == chief_complaint['input_ids'].size(1)-1
        assert len(now_history_tokens) == now_history['input_ids'].size(1)-1
        assert len(past_history_tokens) == past_history['input_ids'].size(1)-1
        assert len(doc_tokens) == doc['input_ids'].size(1)-1

        fields = [chief_complaint,now_history,past_history,doc] # 补充batch维度

        item_entities = set(dic['疾病实体']) | set(dic['治疗实体']) | set(dic['检查实体']) | set(dic['症状实体']) | set(dic['检查结果实体'])

        entities = torch.tensor([0 if entity not in item_entities else 1 for entity in self.opt.entity2id]).unsqueeze(0)

        if torch.cuda.is_available() and self.opt.gpu >= 0:
            for field in fields:
                field['input_ids'] = field['input_ids'].cuda(self.opt.gpu)
                field['attention_mask'] = field['attention_mask'].cuda(self.opt.gpu)
            entities = entities.cuda(self.opt.gpu)

        label = np.array([0 if label not in dic['出院诊断'] else 1 for label in self.opt.label2id])
        with torch.no_grad():
            if not self.use_longformer:
                out,alpha = self.model((chief_complaint['input_ids'], now_history['input_ids'], past_history['input_ids']),
                                   (chief_complaint['attention_mask'], now_history['attention_mask'], past_history['attention_mask']),
                                   entities, self.opt.label_matrix, self.opt.ent_label_matrix)
            else:
                out,alpha = self.model((doc['input_ids'],),(doc['attention_mask'],),entities, self.opt.label_matrix, self.opt.ent_label_matrix)

        out = out.squeeze(0).data.cpu().numpy() # [num_labels]
        alpha = alpha.squeeze(0).data.cpu().numpy() # [num_labels, seq_len]
        alpha = alpha / (alpha.max(axis=1,keepdims=True) + 1e-8)

        data_sizes = []

        for field in fields:
            data_sizes.append(field['input_ids'].size(1)-1)
        visual_data = []
        for labelid in range(self.opt.class_num):
            class_name = self.opt.id2label[labelid]
            cls_type = -1
            if label[labelid] == 1 and out[labelid] >=0: # 模型预测：阳 实际：阳
                cls_type = 0
            elif label[labelid] == 1 and out[labelid] <0: # 模型预测：阴 实际：阳
                cls_type = 1
            elif label[labelid] == 0 and out[labelid] >0: # 模型预测：阳 实际：阴
                cls_type = 2
            else: # 模型预测：阴 实际：阴
                cls_type = 3

            if cls_type == 3:
                continue

            if not self.use_longformer: # alpha : [token1,..,[SEP],token1,...,[SEP],token1,...,[SEP]]
                chief_complaint_attn = alpha[labelid,:data_sizes[0]]
                now_history_attn = alpha[labelid,data_sizes[0]:data_sizes[1]]
                past_history_attn = alpha[labelid,data_sizes[1]:data_sizes[2]]
            else: # longformer
                doc_tokens_tmp = copy.copy(doc_tokens)
                sep_idx1 = doc_tokens_tmp.index(SEP)
                doc_tokens_tmp[sep_idx1] = '#'
                sep_idx2 = doc_tokens_tmp.index(SEP)

                chief_complaint_attn = alpha[labelid,:sep_idx1 + 1]
                now_history_attn = alpha[labelid,sep_idx1+1:sep_idx2+1]
                past_history_attn = alpha[labelid,sep_idx2+1:]


            visual_data.append([{'data':{
                                    '主诉':[chief_complaint_tokens,chief_complaint_attn],\
                                    '现病史':[now_history_tokens,now_history_attn],\
                                    '既往史':[past_history_tokens,past_history_attn]
                                   },
                                 'type':cls_type},class_name])

        visual_data.sort(key = lambda x:x[0]['type'])


        nav_head = ''
        nav_content = ''
        color_id = 0
        for label_item in visual_data:
            labelid = 'disease%d'%self.opt.label2id[label_item[1]]
            nav_head += """
                <li class="nav-item" role="presentation">
                   <a class="nav-link" id="%s-tab" data-toggle="tab" href="#%s" role="tab" aria-controls="%s" aria-selected="false"><span class="label label-%s">%s</span></a>
                </li>
            """%(labelid,labelid,labelid,VISUALIZATION_DICT[label_item[0]['type']],label_item[1])
            nav_content += """
                <div class="tab-pane fade" id="%s" role="tabpanel" aria-labelledby="%s-tab">
                    <font face="楷体">
                        <p>
                            %s
                        </p>
                        <p>
                            %s
                        </p>
                        <p>
                            %s
                        </p>
                    </font>
                </div>
            """%(labelid,labelid,self.field_visual(label_item[0]['data']['主诉'][0],label_item[0]['data']['主诉'][1],COLORS[color_id]),\
                 self.field_visual(label_item[0]['data']['现病史'][0],label_item[0]['data']['现病史'][1],COLORS[color_id]),\
                 self.field_visual(label_item[0]['data']['既往史'][0],label_item[0]['data']['既往史'][1],COLORS[color_id]),)
            color_id += 1
        with open(self.opt.hightlight_path,'r',encoding = 'utf-8') as f:
            html = f.read()

        html = re.sub(r'<ul class=\"nav nav-tabs\" id=\"myTab\" role=\"tablist\">.*?</ul>','<ul class="nav nav-tabs" id="myTab" role="tablist">'+nav_head+'</ul>',html,flags =re.S)
        html = re.sub(r'<div class=\"tab-content\" id=\"myTabContent\" style=\"padding: 10px;font-size: 25px;font-weight:bold;\">.*?</div>','<div class="tab-content" id="myTabContent" style="padding: 10px;font-size: 25px;font-weight:bold;">'+nav_content+'</div>',html,flags =re.S)
        with open(self.opt.output_path,'w',encoding='utf-8') as f:
            f.write(html)

    def field_visual(self,sentence,attn,color):
        span_str = ''
        for c,a in zip(sentence,attn):
            if c in IGNORE_CHAR_SET:
                continue
            span_str += """<span style="background-color:rgb(%d,%d,%d)">%s</span>"""%(255-(255-color[0])*a,\
                                                                255-(255-color[1])*a,255-(255-color[2])*a,c)
        return span_str

